import re
import requests

def print_kv(data, ident=""):
    for key in data:
        print("%s%s: %s" % (ident, key, data.get(key)))


def do_request(p, json=None, data=None):
    print("do_request:", json, data)
    r = None
    if p.method == "POST":
        if json is not None:
            r = requests.post(p.url, headers=p.headers, json=json, verify=False)
        elif data is not None:
            r = requests.post(p.url, headers=p.headers, data=data, verify=False)
    elif p.method == "PUT":
        if json is not None:
            r = requests.put(p.url, headers=p.headers, json=json, verify=False)
        elif data is not None:
            r = requests.put(p.url, headers=p.headers, data=data, verify=False)
    elif p.method == "DELETE":
        if json is not None:
            r = requests.delete(p.url, headers=p.headers, json=json, verify=False)
        elif data is not None:
            r = requests.delete(p.url, headers=p.headers, data=data, verify=False)

    return r


class Parser:
    def __init__(self):
        self.method = ""                # 请求方法
        self.host = ""                  # 主机ip地址或域名
        self.uri = ""                   # 请求uri
        self.url = ""                   # 请求url
        self.body = ""                  # 请求体
        self.proto = "http"             # 请求协议
        self.headers = {}               # 请求头

    # 对请求数据内容进行解析
    def parse(self, data):
        """
        description: 对文件内容进行解析
        arguments:
            data(str): 待解析的数据
        return:
            description: 无
            type: any
        modify_records:
            - 2022-07-23, jetmie, create the action word.
        status: enable
        remarks:
            data获取渠道:
                a) 通过浏览器点击复制/复制为Node.js fetch获取
                b) 在burpsuite中直接拷贝获取
        """
        if data.startswith("fetch(") and data.endswith("});"):
            # 处理拷贝成多个fetch的情况, 对于多个fetch, 只返回第一个
            if data.count('fetch("') > 1:
                _data = data.split('fetch("')[1][:-3]
                return self.parse('fetch("' + _data)

            self.url = re.findall(re.compile(r'fetch\("(.*?)",\s+\{', re.S), data)[0]
            # 对data进行格式化
            if '"body": null,' in data:
                data = data.replace('"body": null,', '"body": "",')

            print(data)
            arguments = eval(re.findall(re.compile(r'(\{.*\})', re.S), data)[0])
            self.headers = arguments.get("headers", {})
            if arguments.get("referrer"):
                self.headers["referrer"] = arguments.get("referrer")

            if arguments.get("body"):
                self.body = arguments.get("body", "")

            self.method = arguments.get("method", "GET")

            sub_str = self.url.split("/")
            self.proto = sub_str[0].strip(":")
            self.host = sub_str[2]
            self.uri = "/" + "/".join(sub_str[3:])
            self.headers["origin"] = f"{self.proto}://{self.host}"

            if self.body and self.headers.get("content-length"):            # 更新content-length字段
                self.headers["content-length"] = len(self.body)
            self.headers["connection"] = "close"
            self.headers["User-Agent"] = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " \
                                         "AppleWebKit/537.36 (KHTML, like Gecko) " \
                                         "Chrome/103.0.5060.114 Safari/537.36 " \
                                         "Edg/103.0.1264.62"
        else:
            if "\n\n" in data:
                data, self.body = data.split("\n\n")[:2]

            self.headers = {}
            for line in data.split("\n"):
                if not line:
                    continue

                if ": " in line:
                    key_index = line.index(": ")
                    key, value = line[:key_index], line[key_index + 2:]
                    self.headers[key.lower()] = value

                    if key.lower() == "host":
                        self.host = value
                else:
                    self.method, self.uri, _ = line.split(" ")

            print(">>>", self.headers, self.host)
            url = self.headers.get("referer") or self.headers.get("origin")
            self.proto = url.split(":")[0] if url else "https"
            self.url = f"{self.proto}://{self.host}{self.uri}"


# 功能: 用于对fetch请求数据进行格式胡
# 参数: data(str): 通过网页Node.js fetch拷贝的内容
# 返回: 格式化后的数据, 数据可以直接拷贝到bp工具使用
def format(data):
    p = Parser()
    p.parse(data)
    ret = ["%s %s HTTP/1.1" % (p.method, p.uri)]

    if "Host" not in p.headers and "host" not in p.headers:
        ret += ["Host: %s" % (p.host, )]
    if "Content-Length" not in p.headers and "content-length" not in p.headers:
        ret += ["Content-Length: %d" % (len(p.body), )]
    for key in p.headers:
        ret += ["%s: %s" % (key, p.headers.get(key))]

    if p.method in ["POST", "PUT", "PATCH", "DELETE"]:
        ret += ["\n" + p.body if p.body else "\n"]
    else:
        ret += ["\n"]
        
    return "\n".join(ret)


# 功能: 通过python发送请求
# 参数: data(str): 通过网页Node.js fetch拷贝的内容
# 返回: 响应结果
def request(data):
    p = Parser()
    p.parse(data)

    if p.method == "GET":
        r = requests.get(p.url, headers=p.headers, verify=False)
    elif p.method in ["POST", "PUT", "DELETE"]:
        if "application/json" in p.headers["content-type"]:
            if p.body.startswith("{") and p.body.endswith("}"):
                p.body = p.body.replace("true", "True")
                p.body = eval(p.body)
                r = requests.post(p.url, headers=p.headers, json=p.body, verify=False)
            elif "&" in p.body and "=" in p.body:
                r = requests.post(p.url, headers=p.headers, data=p.body, verify=False)
        elif "application/x-www-form-urlencoded" in p.headers.get("content-type"):
            r = requests.post(p.url, headers=p.headers, data=p.body, verify=False)
        else:
            return f"Invalid content-type:{p.method}" 
    else:
        return f"Invalid method:{p.method}"

    print("r.headers:", r.headers)
    print("r.content:", r.content)
    ret = [f"code: {r.status_code}"]
    for key in r.headers:
        ret += ["%s: %s" % (key, r.headers.get(key))]
    
    if r.headers.get("content-disposition"):        # 返回为文件的情况
        output = re.findall(re.compile('filename="(.*)"'), r.headers.get("content-disposition"))[0]
        with open(output, "wb") as f:
            f.write(r.content)
        print(f"Save to {output} successful!")
    elif r.headers.get("content-type") and "utf-8" in r.headers.get("content-type").lower():
        ret += ["", r.content.decode("utf-8")]

    return "\n".join(ret)


class Repeater:
    """
    重放类, 对fetch中的内容进行
    """
    def run(self, blob="", verbose=False):
        """
        功能: 执行函数
        参数: blob(str | Parser): 当输入文件时, 则读取并解析文件内容, 将解析后的文件内容发送, 否则
        """
        ret = ""
        if isinstance(blob, Parser):
            p = blob
            p.url = f"{p.proto}://{p.host}{p.uri}"
            if verbose:
                print("request:")
                print("\turl: \n\t\t%s %s" % (p.method, p.url))
                print("\theaders:")
                print_kv(p.headers, "\t\t")
                print("\tbody: \n\t\t%s" % (p.body, ))

            if p.method in ["GET"]:
                r = requests.get(p.url, headers=p.headers, verify=False)
            elif p.method in ["POST", "PUT", "DELETE"]:
                content_type = p.headers.get("content-type", "application/json")
                if "application/json" in content_type:
                    if p.body.startswith("{") and p.body.endswith("}"):
                        p.body = p.body.replace("true", "True")
                        p.body = p.body.replace("false", "False")
                        p.body = p.body.replace("null", "None")
                        print(p.body.encode("utf-8"))
                        p.body = eval(p.body)

                        r = do_request(p, json=p.body)
                    # elif ("&" in p.body and "=" in p.body) or "," in p.body:
                    else:
                        r = do_request(p, data=p.body)
                elif "application/x-www-form-urlencoded" in content_type or "text/plain" in content_type:
                    r = requests.post(p.url, headers=p.headers, data=p.body, verify=False)
                else:
                    print(f"Invalid content-type: {content_type}")
                    return "error"
            else:
                print(f"Invalid method:{p.method}")
                return "error"

            print("r:", r)
            if verbose:
                print("response:")
                print("\tcode: ", r.status_code)
                ret += f"HTTP/1.1 {r.status_code}\n"
                ret += f"Resp-length: {len(r.content)}\n"
                print("\theaders:")
                print(r.headers, type(r.headers))
                for key in r.headers:
                    ret += "%s: %s\n" % (key, r.headers.get(key))
                print_kv(r.headers, "\t\t")

            content_type = r.headers.get("content-type")
            if r.headers.get("content-disposition"):  # 返回为文件的情况
                print(r.headers.get("content-disposition"))
                output = re.findall(re.compile('filename=(.*)'), r.headers.get("content-disposition"))[0]
                output = urllib.parse.unquote(output.strip('"'))
                with open(f"output/{output}", "wb") as f:
                    f.write(r.content)
                print(f"Save to output/{output} successful!")
            else:
                try:
                    resp_body = r.content.decode("utf-8")
                    if verbose:
                        print("\tbody: \n\t\t", resp_body)
                    ret += "\n" + resp_body
                except:
                    if verbose:
                        print("\tbody: \n\t\t", str(r.content))
                    ret += "\n" + str(r.content)

            print("ret:", ret)
            return ret
        elif os.path.isfile(blob):
            p = Parser()
            p.parse_file(blob)

            return self.run(p, True)
        elif isinstance(blob, str):
            p = Parser()
            p.parse(blob)

            return self.run(p, True)
        else:
            print("Invalid input.")


if __name__ == "__main__":
    data = """"""

    # p = Parser()
    # p.parse(data)
    # print(p)
    r = Repeater()
    r.run(data)